import os
import sys
sys.path.append(".")
import cv2
import argparse
import time
import numpy as np
from multiprocessing import Pool
import re
import constant as const
import img_process
import flow
# import compute_flow  # 这一步是另一个工程中，利用深度学习模型计算光流场的
import shutil

def refine_align_section(root_dir, fixed_dir, moving_dir, stage, process_num=-1):
    print(f"root_dir: {root_dir}, fixed_dir: {fixed_dir}, moving_dir: {moving_dir}, stage: {stage}, process_num: {process_num}")
    if stage == 1:
        # 第一步，8nm图片降采样到32nm，每个tile为1024*1024
        img_process.batch_downsample(f"{root_dir}/{fixed_dir}_8nm",
                                    f"{root_dir}/{fixed_dir}_32nm", scale=0.25, process_num=-1)
        img_process.batch_downsample(f"{root_dir}/{moving_dir}_8nm",
                                    f"{root_dir}/{moving_dir}_32nm", scale=0.25, process_num=-1)
        
        # 第二步，32nm图片降采样到256nm，每个tile为128*128
        img_process.batch_downsample(f"{root_dir}/{fixed_dir}_32nm",
                                    f"{root_dir}/{fixed_dir}_256nm", scale=0.125, process_num=-1)
        img_process.batch_downsample(f"{root_dir}/{moving_dir}_32nm",
                                    f"{root_dir}/{moving_dir}_256nm", scale=0.125, process_num=-1)
        
        # # 再将256nm tile小图合并成一张大图(这一步不需要操作)
        # img_process.merge_tiles("/media/hqjin/Elements/em_data/mec/20230321_MEC/wafer52/out/align/refine/11867_256nm",
        #                         "/media/hqjin/Elements/em_data/mec/20230321_MEC/wafer52/out/align/refine/S11867_256nm.png")
        
        # 第三步，将每4个tile合并成1个tile, 即T(i,j), T(i+1,j), T(i,j+1), T(i+1,j+1)合并成1个新的T(i,j)，新的tile为256*256
        img_process.combine_4_tiles_to_larger_tile_with_overlap(f"{root_dir}/{fixed_dir}_256nm",
                                                                f"{root_dir}/{fixed_dir}_256nm_combine")
        img_process.combine_4_tiles_to_larger_tile_with_overlap(f"{root_dir}/{moving_dir}_256nm",
                                                                f"{root_dir}/{moving_dir}_256nm_combine")

        # # 第四步，将第三步生成的两张切片的256*256小图使用深度学习框架进行配准(这一步需要GPU，生成光流场文件)
        # compute_flow.process_image_pairs(f"{root_dir}/{fixed_dir}_256nm_combine",
        #                                  f"{root_dir}/{moving_dir}_256nm_combine",
        #                                  f"{root_dir}/{moving_dir}_256nm_combine_flow")

        shutil.rmtree(f"{root_dir}/{fixed_dir}_256nm")
        shutil.rmtree(f"{root_dir}/{moving_dir}_256nm")
        shutil.rmtree(f"{root_dir}/{fixed_dir}_256nm_combine")
        shutil.rmtree(f"{root_dir}/{moving_dir}_256nm_combine")
    elif stage == 2:
        # 第五步，根据光流场文件夹中每个flow文件(由4块小区域组成，因为第二、三步中每4个tile合并成1个tile，并且tile有重叠), 计算出每个小区域的实际光流场，再合并成一整个大的光流场
        total_flow_256nm = flow.combine_flow(f"{root_dir}/{moving_dir}_256nm_combine_flow", 128)
        
        # 第六步，对整个光流场进行上采样到32nm，并分割为1024*1024的小光流场保存
        flow_32nm = flow.upsample_flow(total_flow_256nm, 8)  # 从256nm上采样到32nm
        flow.split_flow(flow_32nm, 1024, f"{root_dir}/{moving_dir}_32nm_coarse_flow")

        # 第七步，对上一步的小光流场进行合并，并且带overlap，同样的，对32nm图片进行相应的操作
        flow.batch_combine_flows_to_larger_flows_with_overlap(f"{root_dir}/{moving_dir}_32nm_coarse_flow",
                                                            f"{root_dir}/{moving_dir}_32nm_coarse_flow_combine", process_num=-1)
        img_process.batch_combine_tiles_to_larger_tiles_with_overlap(f"{root_dir}/{moving_dir}_32nm",
                                                                    f"{root_dir}/{moving_dir}_32nm_combine", process_num=-1)
        
        # 第八步，对上一步合并的图像应用其对应的光流场得到32nm下align的粗略的结果
        flow.process_flow_to_img(f"{root_dir}/{moving_dir}_32nm_combine",
                                f"{root_dir}/{moving_dir}_32nm_coarse_flow_combine",
                                f"{root_dir}/{moving_dir}_32nm_combine_coarse")
        # # 作为比较，可以将fixed图像也按照第七步合并成多个大图，然后对比第八步的结果
        # img_process.batch_combine_tiles_to_larger_tiles_with_overlap(f"{root_dir}/07246_32nm",
        #                                                              f"{root_dir}/07246_32nm_combine", process_num=-1)

        # 第九步，将上一步的应用光流场得到的图像还原成每一小块不重叠的图像
        img_process.get_small_tiles_from_combine_tiles(f"{root_dir}/{moving_dir}_32nm_combine_coarse",
                                                    f"{root_dir}/{moving_dir}_32nm_coarse",
                                                    1, 52, 1, 54, 1024, const.ROW_OVERLAP, const.COL_OVERLAP)
        # # 可以合并成大图看一下32nm下的结果
        # img_process.merge_tiles("/media/hqjin/Elements/em_data/mec/20230321_MEC/wafer52/out/align/refine/11867_32nm_coarse",
        #                         "/media/hqjin/Elements/em_data/mec/20230321_MEC/wafer52/out/align/refine/11867_32nm_coarse_img.png")

        # 第十步，将上一步的每张32nm下的小图(1024*1024)分割成16个(256*256)的小图
        img_process.split_tiles(f"{root_dir}/{moving_dir}_32nm_coarse",
                                f"{root_dir}/{moving_dir}_32nm_coarse_split", 256)
        img_process.split_tiles(f"{root_dir}/{fixed_dir}_32nm",
                                f"{root_dir}/{fixed_dir}_32nm_split", 256)
        
        # 第十一步，将上一步的每4个tile(256*256)合并成1个tile, 即T(i,j), T(i+1,j), T(i,j+1), T(i+1,j+1)合并成1个新的T(i,j)，新的tile为(512*512)
        img_process.combine_4_tiles_to_larger_tile_with_overlap(f"{root_dir}/{moving_dir}_32nm_coarse_split",
                                                                f"{root_dir}/{moving_dir}_32nm_coarse_split2")
        img_process.combine_4_tiles_to_larger_tile_with_overlap(f"{root_dir}/{fixed_dir}_32nm_split",
                                                                f"{root_dir}/{fixed_dir}_32nm_split2")

        # # # 第十二步，将上一步生成的两张切片的512*512小图使用深度学习框架进行配准(这一步需要GPU，生成光流场文件)
        # compute_flow.process_image_pairs(f"{root_dir}/{fixed_dir}_32nm_split2",
        #                                  f"{root_dir}/{moving_dir}_32nm_coarse_split2",
        #                                  f"{root_dir}/{moving_dir}_32nm_coarse_split2_flow")

        shutil.rmtree(f"{root_dir}/{moving_dir}_256nm_combine_flow")
        shutil.rmtree(f"{root_dir}/{moving_dir}_32nm_coarse_flow")
        shutil.rmtree(f"{root_dir}/{fixed_dir}_32nm")
        shutil.rmtree(f"{root_dir}/{moving_dir}_32nm")
        shutil.rmtree(f"{root_dir}/{moving_dir}_32nm_combine")
        shutil.rmtree(f"{root_dir}/{moving_dir}_32nm_combine_coarse")
        shutil.rmtree(f"{root_dir}/{moving_dir}_32nm_coarse_split")
        shutil.rmtree(f"{root_dir}/{fixed_dir}_32nm_split")
        shutil.rmtree(f"{root_dir}/{fixed_dir}_32nm_split2")
    elif stage == 3:
        # 第十三步，根据光流场文件夹中每个flow文件(由4块小区域组成，因为第十、十一步中每4个tile合并成1个tile，并且tile有重叠), 计算出每个小区域的实际光流场，再合并成一整个大的光流场
        total_flow_32nm = flow.combine_flow(f"{root_dir}/{moving_dir}_32nm_coarse_split2_flow", 256)
        # 再将大的光流场分割为1024*1024的小光流场保存
        flow.split_flow(total_flow_32nm, 1024, f"{root_dir}/{moving_dir}_32nm_fine_flow")

        # 第十四步，对上一步的小光流场进行合并，并且带overlap，同样的，对32nm图片进行相应的操作
        flow.batch_combine_flows_to_larger_flows_with_overlap(f"{root_dir}/{moving_dir}_32nm_fine_flow",
                                                            f"{root_dir}/{moving_dir}_32nm_fine_flow_combine", process_num=-1)
        img_process.batch_combine_tiles_to_larger_tiles_with_overlap(f"{root_dir}/{moving_dir}_32nm_coarse",
                                                                    f"{root_dir}/{moving_dir}_32nm_coarse_combine", process_num=-1)
        
        # 第十五步，对上一步合并的图像应用其对应的光流场得到32nm下align的粗略的结果
        flow.process_flow_to_img(f"{root_dir}/{moving_dir}_32nm_coarse_combine",
                                f"{root_dir}/{moving_dir}_32nm_fine_flow_combine",
                                f"{root_dir}/{moving_dir}_32nm_combine_fine")
        
        # 第十六步，将上一步的应用光流场得到的图像还原成每一小块不重叠的图像（后面检查一下这一步是否没必要）
        img_process.get_small_tiles_from_combine_tiles(f"{root_dir}/{moving_dir}_32nm_combine_fine",
                                                    f"{root_dir}/{moving_dir}_32nm_fine",
                                                    1, 52, 1, 54, 1024, const.ROW_OVERLAP, const.COL_OVERLAP)
        # img_process.merge_tiles(f"{root_dir}/07247_32nm_fine",
        #                         f"{root_dir}/07247_32nm_fine_img.png")

        # 第十七步，将32nm粗略的合并的光流场和精细的合并的光流场做复合运算
        flow.composite_flows(f"{root_dir}/{moving_dir}_32nm_coarse_flow_combine",
                            f"{root_dir}/{moving_dir}_32nm_fine_flow_combine",
                            f"{root_dir}/{moving_dir}_32nm_comp_flow_combine")

        # 第十八步，将32nm复合的合并的光流场拆解为一个个小的光流场(1024*1024*2)
        flow.get_small_flows_from_combine_flows(f"{root_dir}/{moving_dir}_32nm_comp_flow_combine",
                                                f"{root_dir}/{moving_dir}_32nm_comp_flow",
                                                1, 52, 1, 54, 1024, const.ROW_OVERLAP, const.COL_OVERLAP)

        # 第十九步，将上一步的32nm的一个个小光流场(1024*1024*2)进行上采样至8nm(4096*4096*2)
        flow.batch_upsample_flows(f"{root_dir}/{moving_dir}_32nm_comp_flow",
                                f"{root_dir}/{moving_dir}_8nm_flow", 4, 50)
        
        # 第二十步，将8nm下的一个个小光流场进行合并，并且带overlap，同样的，对8nm图片进行相应的操作
        flow.batch_combine_flows_to_larger_flows_with_overlap(f"{root_dir}/{moving_dir}_8nm_flow",
                                                            f"{root_dir}/{moving_dir}_8nm_flow_combine", process_num=50)
        img_process.batch_combine_tiles_to_larger_tiles_with_overlap(f"{root_dir}/{moving_dir}_8nm",
                                                                    f"{root_dir}/{moving_dir}_8nm_combine", process_num=-1)
        
        # 第二十一步，对上一步合并的图像应用其对应的光流场得到8nm下align的粗略的结果
        flow.process_flow_to_img(f"{root_dir}/{moving_dir}_8nm_combine",
                                f"{root_dir}/{moving_dir}_8nm_flow_combine",
                                f"{root_dir}/{moving_dir}_8nm_combine_fine")

        # 第二十二步，将上一步的应用光流场得到的图像还原成每一小块不重叠的图像
        img_process.get_small_tiles_from_combine_tiles(f"{root_dir}/{moving_dir}_8nm_combine_fine",
                                                    f"{root_dir}/{moving_dir}_8nm_fine",
                                                    1, 52, 1, 54, 4096, const.ROW_OVERLAP, const.COL_OVERLAP)
        
        shutil.rmtree(f"{root_dir}/{moving_dir}_32nm_coarse_flow_combine")
        shutil.rmtree(f"{root_dir}/{moving_dir}_32nm_coarse")
        shutil.rmtree(f"{root_dir}/{moving_dir}_32nm_coarse_split2_flow")
        shutil.rmtree(f"{root_dir}/{moving_dir}_32nm_fine_flow")
        shutil.rmtree(f"{root_dir}/{moving_dir}_32nm_fine_flow_combine")
        shutil.rmtree(f"{root_dir}/{moving_dir}_32nm_coarse_combine")
        shutil.rmtree(f"{root_dir}/{moving_dir}_32nm_combine_fine")
        shutil.rmtree(f"{root_dir}/{moving_dir}_32nm_fine")
        shutil.rmtree(f"{root_dir}/{moving_dir}_32nm_comp_flow_combine")
        shutil.rmtree(f"{root_dir}/{moving_dir}_32nm_comp_flow")


def main():
    parser = argparse.ArgumentParser(description="This is a auxiliary tool.")
    parser.add_argument("-r", "--root_dir", type=str, default="/CX/neuro_segment/user/jinhaiqun/out/mec/align_imgs/refine", help="source dir")
    parser.add_argument("-f", "--fixed_dir", type=str, default="01216", help="fixed image dir")
    parser.add_argument("-m", "--moving_dir", type=str, default="01217", help="moving image dir")
    parser.add_argument("-s", "--stage", type=int, default=1, choices=[1, 2, 3], help="moving image dir")
    parser.add_argument("-p", "--process_num", type=int, default=16, help="the number of processes to use(default: 16)")
    args = parser.parse_args()

    refine_align_section(args.root_dir, args.fixed_dir, args.moving_dir, args.stage, 112)


if __name__ == '__main__':
    main()
